# -*- coding: utf-8 -*-

from __future__ import unicode_literals, print_function, division
from io import open
import glob
import unicodedata
import string
import random
import time
import math
from torch.utils.data import DataLoader
from Dataset import TextDataset
from Model import RNN


all_letters = string.printable[:-5]
n_letters = len(all_letters) + 1 # Plus EOS marker
class_num = n_letters # 目标分类数

learning_rate = 0.0002
momentum=0.09
n_iters = 1000

batch_size = 1 # 由于多个sequence放一起的时候需要进行padding，做完padding在grad的时候还需要校验长度，所以最好先不要把batch_size改到1以上。


def inputTransform(line):
    '''
    将字符串转为按字符的one-hot matrix
    '''
    tensor = torch.zeros(len(line), 1, n_letters)
    for letter ,li in enumerate(line):
        tensor[li][0][all_letters.find(letter)] = 1
    return tensor

def targetTransform(line):
    '''
    生成对应的label tensor供NLLLoss使用
    '''
    letter_indexes = [all_letters.find(line[li]) for li in range(1, len(line))]
    letter_indexes.append(n_letters - 1) # EOS
    target_tensor = torch.LongTensor(letter_indexes)
    return target_tensor

textLoader = DataLoader(
        TextDataset(root_dir="./data", transform=inputTransform, target_transform=targetTransform),
        batch_size = batch_size,
        shuffle = True,
    )


criterion = nn.NLLLoss()


def timeSince(since):
    now = time.time()
    s = now - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)

def train():
    model = RNN(n_letters, hidden_size, n_letters)
    hs = model.initHidden()
    optimizer = optim.SGD(model.parameters(), lr = learning_rate, momentum=momentum)

    print_every = n_iters / 100
    start = time.time()

    for ibatch, (input_, target) in enumerate(textLoader):
        optimizer.zero_grad()
        output, hs = model(input_, hs)
        loss = criterion(output, target)

        if ibatch % print_every == 0:
            print('%s (%d %d%%) %.4f' % (timeSince(start), iter, iter / n_iters * 100, loss))

        loss.backward()
        optimizer.step()


if __name__=="__main__":
    train()
